# -*- coding: utf-8 -*-
from PyQt6.QtWidgets import QMessageBox
from collections import defaultdict

class Scheduler:
    def __init__(self, max_students_per_day=8000):
        self.max_students_per_day = max_students_per_day
        self.time_slots = ['第一节', '第二节', '第三节']
        self.days = ['周一', '周二', '周三', '周四', '周五']
    
    def can_schedule_class(self, class_data, day, schedule_table):
        """检查是否可以安排班级到指定日期"""
        current_total = sum(
            sum(c['male'] + c['female'] for c in cell.classes)
            for cell in self.get_day_cells(day, schedule_table)
        )
        return current_total + class_data['male'] + class_data['female'] <= self.max_students_per_day
    
    def get_day_cells(self, day, schedule_table):
        """获取指定日期的所有单元格"""
        col = self.days.index(day)
        return [schedule_table.item(row, col) for row in range(3)]
    
    def group_classes_by_college_grade(self, classes):
        """按学院和年级分组班级"""
        groups = defaultdict(list)
        for class_info in classes:
            key = (class_info['college'], class_info['grade'])
            groups[key].append(class_info)
        return groups
    
    def find_best_slot_for_group(self, group, schedule_table, slot_students, existing_schedules):
        """为同学院同年级的班级组找到最佳时间槽"""
        # 计算组内学生总数
        group_size = sum(c['male'] + c['female'] for c in group)
        college = group[0]['college']
        grade = group[0]['grade']
        
        # 1. 首先查找是否已有同学院同年级的班级被安排
        college_grade_slots = {}
        for (day, slot), classes in existing_schedules.items():
            for class_info in classes:
                if class_info['college'] == college and class_info['grade'] == grade:
                    college_grade_slots[(day, slot)] = True
        
        # 如果找到已安排的同学院同年级时间槽，优先使用这些时间槽
        if college_grade_slots:
            for day, slot in college_grade_slots.keys():
                # 检查该时间槽是否还有足够容量
                daily_total = sum(
                    slot_students.get((day, s), 0)
                    for s in self.time_slots
                )
                if daily_total + group_size <= self.max_students_per_day:
                    return (day, slot), group_size
        
        # 2. 如果没有找到已安排的同学院同年级时间槽，则找最空的时间槽
        min_students = float('inf')
        best_slot = None
        
        for day in self.days:
            for slot in self.time_slots:
                # 计算当天已安排的学生总数
                daily_total = sum(
                    slot_students.get((day, s), 0)
                    for s in self.time_slots
                )
                
                # 检查是否超过每日上限
                if daily_total + group_size > self.max_students_per_day:
                    continue
                
                # 获取当前时间槽的学生数
                current_slot_students = slot_students.get((day, slot), 0)
                
                # 优先选择学生数较少的时间槽
                if current_slot_students < min_students:
                    min_students = current_slot_students
                    best_slot = (day, slot)
        
        return best_slot, group_size
    
    def auto_schedule(self, unscheduled_classes, schedule_table, parent):
        """自动排课功能，优先将同学院同年级的班级安排在同一时间"""
        if not unscheduled_classes:
            QMessageBox.warning(parent, "警告", "没有需要安排的班级")
            return []
        
        # 创建时间槽列表
        time_slots = []
        for day in self.days:
            for slot in self.time_slots:
                time_slots.append((day, slot))
        
        # 获取每个时间槽当前的学生数和已安排的班级
        slot_students = {}
        existing_schedules = {}
        for day, slot in time_slots:
            col = self.days.index(day)
            row = self.time_slots.index(slot)
            cell = schedule_table.item(row, col)
            
            # 记录学生数
            total = sum(c['male'] + c['female'] for c in cell.classes) if cell.classes else 0
            slot_students[(day, slot)] = total
            
            # 记录已安排的班级
            if cell.classes:
                existing_schedules[(day, slot)] = cell.classes.copy()
        
        # 按学院和年级分组
        class_groups = self.group_classes_by_college_grade(unscheduled_classes)
        
        # 对组按总人数从大到小排序（先安排人数多的组）
        sorted_groups = sorted(
            class_groups.values(),
            key=lambda g: sum(c['male'] + c['female'] for c in g),
            reverse=True
        )
        
        scheduled_classes = []
        
        # 首先尝试安排同学院同年级的班级组
        for group in sorted_groups:
            best_slot, group_size = self.find_best_slot_for_group(
                group, schedule_table, slot_students, existing_schedules
            )
            
            if best_slot is None:
                continue  # 没有合适的时间槽，跳过这组
                
            day, slot = best_slot
            col = self.days.index(day)
            row = self.time_slots.index(slot)
            cell = schedule_table.item(row, col)
            
            # 更新课表
            for class_info in group:
                cell.classes.append(class_info)
                scheduled_classes.append(class_info)
                
            # 更新时间槽的学生数
            slot_students[best_slot] = slot_students.get(best_slot, 0) + group_size
            
            # 更新已安排的班级
            if best_slot in existing_schedules:
                existing_schedules[best_slot].extend(group)
            else:
                existing_schedules[best_slot] = group.copy()
        
        # 处理剩余的未安排班级（如果有）
        remaining_classes = [c for c in unscheduled_classes if c not in scheduled_classes]
        if remaining_classes:
            # 对剩余的未安排班级按人数从大到小排序
            remaining_classes.sort(key=lambda x: x['male'] + x['female'], reverse=True)
            
            for class_info in remaining_classes:
                class_size = class_info['male'] + class_info['female']
                best_slot = None
                min_students = float('inf')
                
                # 查找最空的时间槽
                for day in self.days:
                    # 计算当天已安排的学生总数
                    daily_total = sum(
                        slot_students.get((day, s), 0)
                        for s in self.time_slots
                    )
                    
                    # 检查是否超过每日上限
                    if daily_total + class_size > self.max_students_per_day:
                        continue
                    
                    # 检查每个时间槽
                    for slot in self.time_slots:
                        current_slot_students = slot_students.get((day, slot), 0)
                        if current_slot_students < min_students:
                            min_students = current_slot_students
                            best_slot = (day, slot)
                
                if best_slot is None:
                    continue  # 没有合适的时间槽，跳过这个班级
                
                # 更新课表
                day, slot = best_slot
                col = self.days.index(day)
                row = self.time_slots.index(slot)
                cell = schedule_table.item(row, col)
                cell.classes.append(class_info)
                
                # 更新学生数
                slot_students[best_slot] = slot_students.get(best_slot, 0) + class_size
                
                # 添加到已安排列表
                scheduled_classes.append(class_info)
        
        return scheduled_classes